{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eQk5Dv3jyheg"
   },
   "source": [
    "# If you are using Colab for free, we highly recommend you activate the T4 GPU\n",
    "# hardware accelerator. Our models are designed to run with at least 16GB\n",
    "# of RAM, activating T4 will grant the notebook 16GB of GDDR6 RAM as opposed\n",
    "# to the ~13GB Colab gives automatically.\n",
    "# To activate T4:\n",
    "# 1. click on the \"Runtime\" tab\n",
    "# 2. click on \"Change runtime type\"\n",
    "# 3. select T4 GPU under Hardware Accelerator\n",
    "# NOTE: there is a weekly usage limit on using T4 for free"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "POFehtFWxwlz"
   },
   "source": [
    "\n",
    " This example shows an end-to-end recipe for creating a CustomTable, and then creating an Agent process that\n",
    "    will query the table using natural language.\n",
    "\n",
    "    Please note that this example is a 'generalized' and updated version of an earlier example -\n",
    "    \"text2sql-end-to-end-2.py\" - now using the more powerful CustomTables class integrated into the LLMfx process\n",
    "\n",
    "    The example shows the following steps:\n",
    "\n",
    "    1.  Creating a custom table resource from a sample CSV file, included in the slim-sql-tool kit, and also\n",
    "        available in the Examples section with Structured_Tables (customer_table.csv)\n",
    "\n",
    "    2  Asking basic natural language questions:\n",
    "        A.  Looks up the table schema\n",
    "        B.  Packages the table schema with query\n",
    "        C.  Runs inference to convert text into SQL\n",
    "        D.  Queries the database with the generated SQL\n",
    "        E.  Returns result\n",
    "\n",
    "    3.  Using CustomtTable class, this can be run on either Postgres or SQLite DB.\n",
    "\n",
    "    Note: as you substitute for your own CSV and JSON, check out the other examples in this section for loading\n",
    "    configuration ideas and options.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gWyhnhW4v9zm"
   },
   "source": [
    "This notebook is designed for the google collab environment\n",
    "\n",
    "\n",
    "1) The first cell will allow this notebook to access files in your google drive. The example csv for this notebook is located at https://github.com/llmware-ai/llmware/blob/main/examples/Structured_Tables/customer_table.csv\n",
    "\n",
    "Upload it to your drive, and ensure that when you input parameters you point to the right path where the csv is located.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AFH-qyZ_tVpn",
    "outputId": "1a51d6ea-af94-441c-c99b-e2457108e8dc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
     ]
    }
   ],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sDjd_1IPoGx9",
    "outputId": "575d6053-98b6-4098-defb-e1c3cce6777b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: llmware in /usr/local/lib/python3.10/dist-packages (0.3.0)\n",
      "Requirement already satisfied: boto3>=1.24.53 in /usr/local/lib/python3.10/dist-packages (from llmware) (1.34.124)\n",
      "Requirement already satisfied: huggingface-hub>=0.19.4 in /usr/local/lib/python3.10/dist-packages (from llmware) (0.23.2)\n",
      "Requirement already satisfied: numpy>=1.23.2 in /usr/local/lib/python3.10/dist-packages (from llmware) (1.25.2)\n",
      "Requirement already satisfied: pymongo>=4.7.0 in /usr/local/lib/python3.10/dist-packages (from llmware) (4.7.3)\n",
      "Requirement already satisfied: tokenizers>=0.15.0 in /usr/local/lib/python3.10/dist-packages (from llmware) (0.19.1)\n",
      "Requirement already satisfied: psycopg-binary==3.1.17 in /usr/local/lib/python3.10/dist-packages (from llmware) (3.1.17)\n",
      "Requirement already satisfied: psycopg==3.1.17 in /usr/local/lib/python3.10/dist-packages (from llmware) (3.1.17)\n",
      "Requirement already satisfied: pgvector==0.2.4 in /usr/local/lib/python3.10/dist-packages (from llmware) (0.2.4)\n",
      "Requirement already satisfied: colorama==0.4.6 in /usr/local/lib/python3.10/dist-packages (from llmware) (0.4.6)\n",
      "Requirement already satisfied: librosa>=0.10.0 in /usr/local/lib/python3.10/dist-packages (from llmware) (0.10.2.post1)\n",
      "Requirement already satisfied: typing-extensions>=4.1 in /usr/local/lib/python3.10/dist-packages (from psycopg==3.1.17->llmware) (4.12.1)\n",
      "Requirement already satisfied: botocore<1.35.0,>=1.34.124 in /usr/local/lib/python3.10/dist-packages (from boto3>=1.24.53->llmware) (1.34.124)\n",
      "Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from boto3>=1.24.53->llmware) (1.0.1)\n",
      "Requirement already satisfied: s3transfer<0.11.0,>=0.10.0 in /usr/local/lib/python3.10/dist-packages (from boto3>=1.24.53->llmware) (0.10.1)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->llmware) (3.14.0)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->llmware) (2023.6.0)\n",
      "Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->llmware) (24.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->llmware) (6.0.1)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->llmware) (2.31.0)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->llmware) (4.66.4)\n",
      "Requirement already satisfied: audioread>=2.1.9 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (3.0.1)\n",
      "Requirement already satisfied: scipy>=1.2.0 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (1.11.4)\n",
      "Requirement already satisfied: scikit-learn>=0.20.0 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (1.2.2)\n",
      "Requirement already satisfied: joblib>=0.14 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (1.4.2)\n",
      "Requirement already satisfied: decorator>=4.3.0 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (4.4.2)\n",
      "Requirement already satisfied: numba>=0.51.0 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (0.58.1)\n",
      "Requirement already satisfied: soundfile>=0.12.1 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (0.12.1)\n",
      "Requirement already satisfied: pooch>=1.1 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (1.8.1)\n",
      "Requirement already satisfied: soxr>=0.3.2 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (0.3.7)\n",
      "Requirement already satisfied: lazy-loader>=0.1 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (0.4)\n",
      "Requirement already satisfied: msgpack>=1.0 in /usr/local/lib/python3.10/dist-packages (from librosa>=0.10.0->llmware) (1.0.8)\n",
      "Requirement already satisfied: dnspython<3.0.0,>=1.16.0 in /usr/local/lib/python3.10/dist-packages (from pymongo>=4.7.0->llmware) (2.6.1)\n",
      "Requirement already satisfied: python-dateutil<3.0.0,>=2.1 in /usr/local/lib/python3.10/dist-packages (from botocore<1.35.0,>=1.34.124->boto3>=1.24.53->llmware) (2.8.2)\n",
      "Requirement already satisfied: urllib3!=2.2.0,<3,>=1.25.4 in /usr/local/lib/python3.10/dist-packages (from botocore<1.35.0,>=1.34.124->boto3>=1.24.53->llmware) (2.0.7)\n",
      "Requirement already satisfied: llvmlite<0.42,>=0.41.0dev0 in /usr/local/lib/python3.10/dist-packages (from numba>=0.51.0->librosa>=0.10.0->llmware) (0.41.1)\n",
      "Requirement already satisfied: platformdirs>=2.5.0 in /usr/local/lib/python3.10/dist-packages (from pooch>=1.1->librosa>=0.10.0->llmware) (4.2.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.19.4->llmware) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.19.4->llmware) (3.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.19.4->llmware) (2024.6.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn>=0.20.0->librosa>=0.10.0->llmware) (3.5.0)\n",
      "Requirement already satisfied: cffi>=1.0 in /usr/local/lib/python3.10/dist-packages (from soundfile>=0.12.1->librosa>=0.10.0->llmware) (1.16.0)\n",
      "Requirement already satisfied: pycparser in /usr/local/lib/python3.10/dist-packages (from cffi>=1.0->soundfile>=0.12.1->librosa>=0.10.0->llmware) (2.22)\n",
      "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil<3.0.0,>=2.1->botocore<1.35.0,>=1.34.124->boto3>=1.24.53->llmware) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install llmware"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "zAUBmgSptRgQ"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from llmware.agents import LLMfx\n",
    "from llmware.resources import CustomTable\n",
    "from llmware.configs import LLMWareConfig\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "id": "P0okOaH0taa0"
   },
   "outputs": [],
   "source": [
    "def build_table(db=None, table_name=None, load_fp=None, load_file=None):\n",
    "    \"\"\"Simple example script to take a CSV or JSON/JSONL and create a DB Table.\"\"\"\n",
    "    custom_table = CustomTable(db=db, table_name=table_name)\n",
    "    analysis = custom_table.validate_csv(load_fp, load_file)\n",
    "    print(\"update: analysis from validate_csv: \", analysis)\n",
    "\n",
    "    if load_file.endswith(\".csv\"):\n",
    "        output = custom_table.load_csv(load_fp, load_file)\n",
    "    elif load_file.endswith(\".jsonl\") or load_file.endswith(\".json\"):\n",
    "        output = custom_table.load_json(load_fp, load_file)\n",
    "    else:\n",
    "        print(\"file type not supported for db load\")\n",
    "        return -1\n",
    "\n",
    "    print(\"update: output from loading file: \", output)\n",
    "\n",
    "    sample_range = min(10, len(custom_table.rows))\n",
    "    for x in range(0, sample_range):\n",
    "        print(\"update: sample rows: \", x, custom_table.rows[x])\n",
    "\n",
    "    updated_schema = custom_table.test_and_remediate_schema(samples=20, auto_remediate=True)\n",
    "    print(\"update: updated schema: \", updated_schema)\n",
    "\n",
    "    custom_table.insert_rows()\n",
    "    return 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "id": "Q6uD3D-wtfOG"
   },
   "outputs": [],
   "source": [
    "def agent_natural_language_sql_query(query_list, db=None, table_name=None):\n",
    "    \"\"\"Query a CustomTable in natural language.\"\"\"\n",
    "    agent = LLMfx()\n",
    "    agent.load_tool(\"sql\", sample=False, get_logits=True, temperature=0.0)\n",
    "\n",
    "    for i, query in enumerate(query_list):\n",
    "        response = agent.query_custom_table(query, db=db, table=table_name)\n",
    "\n",
    "    for x in range(len(agent.research_list)):\n",
    "        print(\"research: \", x, agent.research_list[x])\n",
    "\n",
    "    return agent.research_list\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bJaDKQyAthUO",
    "outputId": "1931b424-b49f-4176-8c0f-aa2c68d5304a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "update: analysis from validate_csv:  {'rows': 20, 'columns': 6, 'conforming_rows_percent': 1.0, 'column_frequency_analysis': {6: 20}, 'nonconforming_rows': []}\n",
      "update: output from loading file:  {'rows': 19, 'columns': 6, 'schema': {'customer_name': 'text', 'account_number': 'integer', 'customer_level': 'text', 'vip_customer': 'text', 'annual_spend': 'integer', 'user_name': 'text'}, 'skipped_rows': []}\n",
      "update: sample rows:  0 {'customer_name': 'Martha Williams', 'account_number': '98320893', 'customer_level': 'gold', 'vip_customer': 'yes', 'annual_spend': '63250', 'user_name': 'mwilliams'}\n",
      "update: sample rows:  1 {'customer_name': 'Susan Soinsin', 'account_number': '53439382', 'customer_level': 'silver', 'vip_customer': 'no', 'annual_spend': '112 ', 'user_name': 'ssinsin'}\n",
      "update: sample rows:  2 {'customer_name': 'Michael Rogers', 'account_number': '88888444', 'customer_level': 'bronze', 'vip_customer': 'no', 'annual_spend': '3 ', 'user_name': 'rogersm'}\n",
      "update: sample rows:  3 {'customer_name': 'John Jones', 'account_number': '9898482', 'customer_level': 'gold', 'vip_customer': 'yes', 'annual_spend': '12430', 'user_name': 'jjns'}\n",
      "update: sample rows:  4 {'customer_name': 'Melinda Lyons', 'account_number': '1234953', 'customer_level': 'silver', 'vip_customer': 'no', 'annual_spend': '234', 'user_name': 'mlyons'}\n",
      "update: sample rows:  5 {'customer_name': 'Arvind Arora', 'account_number': '8998899', 'customer_level': 'gold', 'vip_customer': 'yes', 'annual_spend': '42650', 'user_name': 'aarora'}\n",
      "update: sample rows:  6 {'customer_name': 'Paul Rinno', 'account_number': '3829235', 'customer_level': 'silver', 'vip_customer': 'no', 'annual_spend': '1293', 'user_name': 'prinno'}\n",
      "update: sample rows:  7 {'customer_name': 'Vinod Aggarwal', 'account_number': '9389532', 'customer_level': 'platinum', 'vip_customer': 'yes', 'annual_spend': '93540', 'user_name': 'aggarwal'}\n",
      "update: sample rows:  8 {'customer_name': 'Bryan Lewis', 'account_number': '12399853', 'customer_level': 'gold', 'vip_customer': 'no', 'annual_spend': '9732', 'user_name': 'brlewis'}\n",
      "update: sample rows:  9 {'customer_name': 'Allison Winters', 'account_number': '83902867', 'customer_level': 'gold', 'vip_customer': 'yes', 'annual_spend': '15000', 'user_name': 'wintersall'}\n",
      "update: updated schema:  {'customer_name': 'text', 'account_number': 'integer', 'customer_level': 'text', 'vip_customer': 'text', 'annual_spend': 'integer', 'user_name': 'text'}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Input parameters\n",
    "db = \"sqlite\"  # SQLite database path\n",
    "table_name = \"customer_table\"\n",
    "\n",
    "# Path to the CSV file in Google Drive\n",
    "input_fp = \"/content/drive/My Drive\"\n",
    "input_fn = \"customer_table.csv\"\n",
    "\n",
    "# Build the table\n",
    "build_table(db=db, table_name=table_name, load_fp=input_fp, load_file=input_fn)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "FJQlgHC9uq8Q",
    "outputId": "2a9f6715-2cd7-41b5-aca6-ed3c9f05fb0e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "update: Launching LLMfx process\n",
      "step - \t1 - \tcreating object - ready to start processing.\n",
      "step - \t2 - \tloading tool - sql\n",
      "step - \t3 - \tloading new processing text - 1 new entries\n",
      "step - \t4 - \texecuting function call - deploying - text-to-sql\n",
      "\t\t\t\t -- query - Which customers have vip customer status of yes?\n",
      "\t\t\t\t -- table_schema - CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )\n",
      "step - \t5 - \texecuting function call - getting response - sql\n",
      "\t\t\t\t -- llm_response - SELECT customer_name FROM customer_table WHERE vip_customer = 'yes'\n",
      "\t\t\t\t -- output type - text\n",
      "\t\t\t\t -- usage - {'input': 58, 'output': 17, 'total': 75, 'metric': 'tokens', 'processing_time': 1.0267748832702637}\n",
      "step - \t6 - \texecuting research call - executing query on db\n",
      "\t\t\t\t -- db - sqlite\n",
      "\t\t\t\t -- sql_query - SELECT customer_name FROM customer_table WHERE vip_customer = 'yes'\n",
      "step - \t7 - \texecuting research  - getting response - sql\n",
      "\t\t\t\t -- result - [('Martha Williams',), ('John Jones',), ('Arvind Arora',), ('Vinod Aggarwal',), ('Allison Winters',), ('Olivia Smith',), ('Alan Wang',), ('Wilmer Rodriguez',), ('Eliza Listron',), ('Douglas Hudson',), ('Martha Williams',), ('John Jones',), ('Arvind Arora',), ('Vinod Aggarwal',), ('Allison Winters',), ('Olivia Smith',), ('Alan Wang',), ('Wilmer Rodriguez',), ('Eliza Listron',), ('Douglas Hudson',)]\n",
      "step - \t8 - \tloading new processing text - 1 new entries\n",
      "step - \t9 - \texecuting function call - deploying - text-to-sql\n",
      "\t\t\t\t -- query - What is the highest annual spend of any customer?\n",
      "\t\t\t\t -- table_schema - CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )\n",
      "step - \t10 - \texecuting function call - getting response - sql\n",
      "\t\t\t\t -- llm_response - SELECT MAX(annual_spend) FROM customer_table\n",
      "\t\t\t\t -- output type - text\n",
      "\t\t\t\t -- usage - {'input': 57, 'output': 13, 'total': 70, 'metric': 'tokens', 'processing_time': 1.495262861251831}\n",
      "step - \t11 - \texecuting research call - executing query on db\n",
      "\t\t\t\t -- db - sqlite\n",
      "\t\t\t\t -- sql_query - SELECT MAX(annual_spend) FROM customer_table\n",
      "step - \t12 - \texecuting research  - getting response - sql\n",
      "\t\t\t\t -- result - [(93540,)]\n",
      "step - \t13 - \tloading new processing text - 1 new entries\n",
      "step - \t14 - \texecuting function call - deploying - text-to-sql\n",
      "\t\t\t\t -- query - Which customer has account number 1234953\n",
      "\t\t\t\t -- table_schema - CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )\n",
      "step - \t15 - \texecuting function call - getting response - sql\n",
      "\t\t\t\t -- llm_response - SELECT customer_name FROM customer_table WHERE account_number = 1234953\n",
      "\t\t\t\t -- output type - text\n",
      "\t\t\t\t -- usage - {'input': 61, 'output': 21, 'total': 82, 'metric': 'tokens', 'processing_time': 2.2324819564819336}\n",
      "step - \t16 - \texecuting research call - executing query on db\n",
      "\t\t\t\t -- db - sqlite\n",
      "\t\t\t\t -- sql_query - SELECT customer_name FROM customer_table WHERE account_number = 1234953\n",
      "step - \t17 - \texecuting research  - getting response - sql\n",
      "\t\t\t\t -- result - [('Melinda Lyons',), ('Melinda Lyons',)]\n",
      "step - \t18 - \tloading new processing text - 1 new entries\n",
      "step - \t19 - \texecuting function call - deploying - text-to-sql\n",
      "\t\t\t\t -- query - Which customer has the lowest annual spend?\n",
      "\t\t\t\t -- table_schema - CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )\n",
      "step - \t20 - \texecuting function call - getting response - sql\n",
      "\t\t\t\t -- llm_response - SELECT customer_name FROM customer_table ORDER BY annual_spend LIMIT 1\n",
      "\t\t\t\t -- output type - text\n",
      "\t\t\t\t -- usage - {'input': 56, 'output': 17, 'total': 73, 'metric': 'tokens', 'processing_time': 1.4454705715179443}\n",
      "step - \t21 - \texecuting research call - executing query on db\n",
      "\t\t\t\t -- db - sqlite\n",
      "\t\t\t\t -- sql_query - SELECT customer_name FROM customer_table ORDER BY annual_spend LIMIT 1\n",
      "step - \t22 - \texecuting research  - getting response - sql\n",
      "\t\t\t\t -- result - [('Michael Rogers',)]\n",
      "step - \t23 - \tloading new processing text - 1 new entries\n",
      "step - \t24 - \texecuting function call - deploying - text-to-sql\n",
      "\t\t\t\t -- query - Is Susan Soinsin a vip customer?\n",
      "\t\t\t\t -- table_schema - CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )\n",
      "step - \t25 - \texecuting function call - getting response - sql\n",
      "\t\t\t\t -- llm_response - SELECT vip_customer FROM customer_table WHERE customer_name = 'Susan Soinsin'\n",
      "\t\t\t\t -- output type - text\n",
      "\t\t\t\t -- usage - {'input': 57, 'output': 22, 'total': 79, 'metric': 'tokens', 'processing_time': 1.4810450077056885}\n",
      "step - \t26 - \texecuting research call - executing query on db\n",
      "\t\t\t\t -- db - sqlite\n",
      "\t\t\t\t -- sql_query - SELECT vip_customer FROM customer_table WHERE customer_name = 'Susan Soinsin'\n",
      "step - \t27 - \texecuting research  - getting response - sql\n",
      "\t\t\t\t -- result - [('no',), ('no',)]\n",
      "research:  0 {'step': 6, 'tool': 'sql', 'db_response': [('Martha Williams',), ('John Jones',), ('Arvind Arora',), ('Vinod Aggarwal',), ('Allison Winters',), ('Olivia Smith',), ('Alan Wang',), ('Wilmer Rodriguez',), ('Eliza Listron',), ('Douglas Hudson',), ('Martha Williams',), ('John Jones',), ('Arvind Arora',), ('Vinod Aggarwal',), ('Allison Winters',), ('Olivia Smith',), ('Alan Wang',), ('Wilmer Rodriguez',), ('Eliza Listron',), ('Douglas Hudson',)], 'sql_query': \"SELECT customer_name FROM customer_table WHERE vip_customer = 'yes'\", 'query': 'Which customers have vip customer status of yes?', 'db': 'sqlite', 'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'}\n",
      "research:  1 {'step': 11, 'tool': 'sql', 'db_response': [(93540,)], 'sql_query': 'SELECT MAX(annual_spend) FROM customer_table', 'query': 'What is the highest annual spend of any customer?', 'db': 'sqlite', 'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'}\n",
      "research:  2 {'step': 16, 'tool': 'sql', 'db_response': [('Melinda Lyons',), ('Melinda Lyons',)], 'sql_query': 'SELECT customer_name FROM customer_table WHERE account_number = 1234953', 'query': 'Which customer has account number 1234953', 'db': 'sqlite', 'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'}\n",
      "research:  3 {'step': 21, 'tool': 'sql', 'db_response': [('Michael Rogers',)], 'sql_query': 'SELECT customer_name FROM customer_table ORDER BY annual_spend LIMIT 1', 'query': 'Which customer has the lowest annual spend?', 'db': 'sqlite', 'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'}\n",
      "research:  4 {'step': 26, 'tool': 'sql', 'db_response': [('no',), ('no',)], 'sql_query': \"SELECT vip_customer FROM customer_table WHERE customer_name = 'Susan Soinsin'\", 'query': 'Is Susan Soinsin a vip customer?', 'db': 'sqlite', 'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'step': 6,\n",
       "  'tool': 'sql',\n",
       "  'db_response': [('Martha Williams',),\n",
       "   ('John Jones',),\n",
       "   ('Arvind Arora',),\n",
       "   ('Vinod Aggarwal',),\n",
       "   ('Allison Winters',),\n",
       "   ('Olivia Smith',),\n",
       "   ('Alan Wang',),\n",
       "   ('Wilmer Rodriguez',),\n",
       "   ('Eliza Listron',),\n",
       "   ('Douglas Hudson',),\n",
       "   ('Martha Williams',),\n",
       "   ('John Jones',),\n",
       "   ('Arvind Arora',),\n",
       "   ('Vinod Aggarwal',),\n",
       "   ('Allison Winters',),\n",
       "   ('Olivia Smith',),\n",
       "   ('Alan Wang',),\n",
       "   ('Wilmer Rodriguez',),\n",
       "   ('Eliza Listron',),\n",
       "   ('Douglas Hudson',)],\n",
       "  'sql_query': \"SELECT customer_name FROM customer_table WHERE vip_customer = 'yes'\",\n",
       "  'query': 'Which customers have vip customer status of yes?',\n",
       "  'db': 'sqlite',\n",
       "  'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'},\n",
       " {'step': 11,\n",
       "  'tool': 'sql',\n",
       "  'db_response': [(93540,)],\n",
       "  'sql_query': 'SELECT MAX(annual_spend) FROM customer_table',\n",
       "  'query': 'What is the highest annual spend of any customer?',\n",
       "  'db': 'sqlite',\n",
       "  'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'},\n",
       " {'step': 16,\n",
       "  'tool': 'sql',\n",
       "  'db_response': [('Melinda Lyons',), ('Melinda Lyons',)],\n",
       "  'sql_query': 'SELECT customer_name FROM customer_table WHERE account_number = 1234953',\n",
       "  'query': 'Which customer has account number 1234953',\n",
       "  'db': 'sqlite',\n",
       "  'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'},\n",
       " {'step': 21,\n",
       "  'tool': 'sql',\n",
       "  'db_response': [('Michael Rogers',)],\n",
       "  'sql_query': 'SELECT customer_name FROM customer_table ORDER BY annual_spend LIMIT 1',\n",
       "  'query': 'Which customer has the lowest annual spend?',\n",
       "  'db': 'sqlite',\n",
       "  'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'},\n",
       " {'step': 26,\n",
       "  'tool': 'sql',\n",
       "  'db_response': [('no',), ('no',)],\n",
       "  'sql_query': \"SELECT vip_customer FROM customer_table WHERE customer_name = 'Susan Soinsin'\",\n",
       "  'query': 'Is Susan Soinsin a vip customer?',\n",
       "  'db': 'sqlite',\n",
       "  'work_item': 'CREATE TABLE customer_table (customer_name text, account_number integer, customer_level text, vip_customer text, annual_spend integer, user_name text )'}]"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_list = [\n",
    "    \"Which customers have vip customer status of yes?\",\n",
    "    \"What is the highest annual spend of any customer?\",\n",
    "    \"Which customer has account number 1234953\",\n",
    "    \"Which customer has the lowest annual spend?\",\n",
    "    \"Is Susan Soinsin a vip customer?\"\n",
    "]\n",
    "\n",
    "# Execute the queries\n",
    "agent_natural_language_sql_query(query_list, db=db, table_name=table_name)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
